Add inter-document attention masking to GPTDataset - #5298
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
bbfb846 to
f0d55be
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| def get_batch_on_this_tp_rank( | ||
| batch: dict[str, torch.Tensor], | ||
| is_sft: bool, | ||
| is_packed: bool, |
There was a problem hiding this comment.
The parameter rename from is_sft to is_packed wasn't propagated to megatron/elastification/pretrain_hybrid_flex.py, which still calls get_batch_on_this_tp_rank(..., is_sft=is_sft, ...) at line 205. This will raise a TypeError at runtime.
There was a problem hiding this comment.
Two issues:
-
Bug — padding not reflected in
cu_seqlens(gpt_dataset.py):document_lengthsis computed before padding is appended tosample_parts, so when a sample is shorter thansequence_length,cu_seqlens[-1] < sequence_lengthwhiletokens.shape[0] == sequence_length. FlashAttention varlen will either error or silently ignore tail positions. -
Broken caller (
elastification/pretrain_hybrid_flex.py): Still passesis_sft=is_sfttoget_batch_on_this_tp_rankafter the parameter was renamed tois_packed. This will raise aTypeErrorat runtime.
|
/claude review |
|
/claude review |
deepakn94
changed the title
|
|
||
| elif is_sft: | ||
| elif is_packed: | ||
| # NOTE(asolergi-nv): Broadcast required THD metadata for SFT to intermediate stages |
There was a problem hiding this comment.
Nit: this NOTE (and the identical one at the other elif is_packed: branch below) still says "SFT" but now also applies to --dataloader-pack-sequences. Worth updating to mention packed sequences generally.
There was a problem hiding this comment.
Clean PR — the rename from is_sft to is_packed is thorough and the new tests cover the important configurations well.
One potential correctness issue flagged inline: when a sample ends exactly at a document boundary, document_lengths[-1] -= 1 can produce a zero-length segment in cu_seqlens. See the inline comment on gpt_dataset.py for details and a suggested fix.
|
/claude review |
6766317 to
1fd9d11
Compare
|
/claude review |
There was a problem hiding this comment.
Clean PR overall — the is_sft → is_packed rename is thorough, tests cover both the dataset and TP-broadcast paths, and the argument validation is in place.
One bug flagged inline: a zero-length cu_seqlens segment can occur when a sample boundary aligns exactly with a document boundary (the C helper sets doc_index_end_offset = 0, the extra-token-only part becomes length 0 after subtraction). See the inline comment for a suggested fix.
631247e to
897e6ac
Compare
Compute cu_seqlens and max_seqlen from document boundaries within each sample, and reset position IDs per document. This lets FlashAttention / TransformerEngine (and linear attention) restrict attention to individual documents via the varlen (thd) path, replacing the legacy dense attention mask built by --reset-attention-mask. Also updates FIMDataset._query_document_sample_shuffle_indices to return the 3-tuple (sample, document_ids, document_lengths) expected by the new code path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
Add use_per_sequence_balancing parameter to get_batch_on_this_cp_rank so inter-document masking routes through per-sequence zigzag partitioning (document lengths are arbitrary and not divisible by 2 * cp_size). Remove the CP > 1 assertion in arguments.py and pass the new flag from all three pretrain entry points. Extend test_inter_document_masking_batch to parametrize over cp_size and add test_get_batch_on_this_cp_rank_per_sequence_balancing to verify zigzag chunk selection with known token values. Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
|
/claude review |
There was a problem hiding this comment.
Light review — one bug found in the test, posted as an inline comment.
The production code changes look correct: the is_sft → has_cu_seqlens rename is consistently applied, the get_batch_on_this_cp_rank refactoring preserves existing behavior when use_per_sequence_balancing is False (the default), and the cu_seqlens/position-ID computation in __getitem__ handles the edge cases (empty last document after extra-token subtraction, padding shortfall) properly. No new direct process group access in megatron/core production code.
Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
|
/claude review |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28275713645 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28276368095 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28317761167 |
|
Note: context parallelism ( |
6 participants
Summary
inter_document_maskingconfig option toGPTDatasetConfigthat computescu_seqlensandmax_seqlenfrom document boundaries within each sample, and resets position IDs per document.thd) path, replacing the legacy dense attention mask approach used by--reset-attention-mask.cu_seqlens/max_seqlenthrough TP broadcast (get_batch_on_this_tp_rank) so non-zero TP ranks receive the metadata.--dataloader-inter-document-maskingCLI argument.FIMDataset._query_document_sample_shuffle_indicesto return the 3-tuple expected by the new code path._get_batch_on_this_cp_rank_per_document_balancing).cu_seqlensto a fixed length inGPTDatasetsodefault_collatecan stack samples with different numbers of documents whenmicro_batch_size > 1.Results
8B-scale training runs (300B tokens) comparing baseline vs inter-document masking for both Transformer and Hybrid (Mamba+attention) architectures. Inter-document masking produces lower training and validation loss for both architectures, with no throughput regression. Gradient norms and parameter norms track closely between baseline and masking variants.
Test plan
test_inter_document_masking) verifyingcu_seqlensvalues, position ID resets, andmax_seqlenacross 20 samples.test_inter_document_masking_batch) verifying end-to-endget_batchwith inter-document masking across TP and PP ranks (parametrized over TP={1,2,4} x PP={1,2,4}).test_hybrid_cp_batch) verifying CP partitioning with inter-document masking across TP, PP, and CP ranks.test_flatten_batch_for_packed_sequences_padded_cu_seqlens(on main via Merge cu_seqlens across micro-batch for THD attention #5454).--dataloader-inter-document-maskingcomparing loss curves against baseline (see Results above).